from aoct import Aoct_Model
from os.path import join, dirname
import os
import librosa
from reader import load_audio
import datetime
import Globaldata
import torch
error_threld = 0.3
def test(file_name='assets/test.wav', model_path='assets/2024-6-29-22121.pth',device='cpu'):
    device = torch.device(device=device)
    file_name = join(dirname(__file__), file_name)
    #for test
    # model = Aoct_Model()
    # model.to("cpu")
    # model.load_state_dict(torch.load(model_path, map_location="cpu"))
    model = Globaldata.initmodel # for android
    model.eval()
    audio, fs = librosa.load(file_name, sr=1000)
    audios = load_audio(audio,fs)
    wav_len = len(audios)
    print("this audio lens is :", wav_len)
    i = 0
    start = datetime.datetime.now()
    for audio in audios:
        audio = torch.from_numpy(audio)
        if device:
            audio = audio.to(device)
        with torch.no_grad():
            audio = audio.to(device)
            audio = torch.unsqueeze(audio, dim=0)
            output = model(audio)
            output_prob = torch.sigmoid(output)  # 使用 Sigmoid 获取概率
            prediction = (output_prob > error_threld).float()
        prediction_np = prediction.cpu().numpy()
        if prediction_np == 1:
            i = i + 1
            a = int(wav_len * 0.2)
            if i == max(a,1):
                end = datetime.datetime.now()
                return "Abnorm-"+str((end-start))
                # return 1
    return "Norm"
    # return 0


def test_for_foloer(root_path):
    files = os.listdir(root_path)
    count_0 = 0
    count_1 = 0
    total_count = len(files)
    i = 1000
    total_count = i if i < total_count else total_count
    for file in files:
        i=i-1
        if i < 0:
            break
        output = test(
            file_name=os.path.join(root_path, file),
        )
        if output == 0:
            count_0 += 1
        elif output == 1:
            count_1 += 1

    percent_0 = count_0 / total_count * 100
    percent_1 = count_1 / total_count * 100

    print('类别0数量:', count_0)
    print('类别1数量:', count_1)
    print('类别0所占百分比:', percent_0)
    print('类别1所占百分比:', percent_1)
# for folder
if __name__ == '__main__':
    # print("==============================unhealth==============================")
    # root_path = './assets/health'
    # test_for_foloer(root_path)
    print("==============================unhealth==============================")
    root_path = './assets/health'
    root_path = './assets/unhealthy1'
    test_for_foloer(root_path)


# if __name__ == '__main__':
#     result = test("./assets/unhealthy/1718243686932.wav")
#     print(result)